package com.smh.aiagent.chatmemory;

import com.esotericsoftware.kryo.Kryo;
import com.esotericsoftware.kryo.io.Input;
import com.esotericsoftware.kryo.io.Output;
import org.objenesis.strategy.StdInstantiatorStrategy;
import org.springframework.ai.chat.memory.ChatMemory;
import org.springframework.ai.chat.messages.Message;

import java.io.File;
import java.io.FileInputStream;
import java.io.FileOutputStream;
import java.util.ArrayList;
import java.util.List;

/**
 * @author shiminghui
 * @date 2025/5/31 16:05
 * @description: TODO
 */
public class FileBasedChatMemory implements ChatMemory {


    private static final Kryo kryo = new Kryo();

    static {
        kryo.setRegistrationRequired(false); // 允许未注册的类
        kryo.setInstantiatorStrategy(new StdInstantiatorStrategy()); // 使用默认的实例化策略
    }

    private final String BASE_DIR;

    public FileBasedChatMemory(String baseDir) {
        this.BASE_DIR = baseDir;
        File dir = new File(BASE_DIR);
        if (!dir.exists()) {
            dir.mkdirs();
        }
    }


    @Override
    public void add(String conversationId, Message message) {
        List<Message> messageList = getOrCreateConversation(conversationId);
        messageList.add(message);
        saveToFile(conversationId, messageList);
    }

    @Override
    public void add(String conversationId, List<Message> messages) {
        List<Message> messageList = getOrCreateConversation(conversationId);
        messageList.addAll(messages);
        saveToFile(conversationId, messageList);
    }

    @Override
    public List<Message> get(String conversationId, int lastN) {
        List<Message> messageList = getOrCreateConversation(conversationId);
        return messageList.stream().skip(Math.max(0, messageList.size() - lastN)).toList();
    }

    @Override
    public void clear(String conversationId) {
        File conversationFile = getConversationFile(conversationId);
        if (conversationFile.exists()) {
            conversationFile.delete();
        }
    }

    public void saveToFile(String conversationId, List<Message> messages) {
        File file = getConversationFile(conversationId);
        try (Output output = new Output(new FileOutputStream(file))) {
            kryo.writeObject(output, messages);
        } catch (Exception e) {
            e.printStackTrace();
        }
    }

    public List<Message> getOrCreateConversation(String conversationId) {
        File file = getConversationFile(conversationId);
        List<Message> messageList = new ArrayList<>();
        if (file.exists()) {
            try (Input input = new Input(new FileInputStream(file));) {
                return kryo.readObject(input, ArrayList.class);
            } catch (Exception e) {
                e.printStackTrace();
            }
        }
        return messageList;
    }


    public File getConversationFile(String conversationId) {
        return new File(BASE_DIR, conversationId + ".kryo");
    }

}
